import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.tree import DecisionTreeClassifier
from sklearn.model_selection import train_test_split
from sklearn.datasets import load_wine
import matplotlib.pyplot as plt
import datetime

# 创建日志文件
def create_log_file():
    current_time = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    log_file = f"random_forest_experiment_{current_time}.log"
    with open(log_file, 'w') as f:
        f.write(f"实验时间: {datetime.datetime.now()}\n")
        f.write("=" * 75 + "\n")
    return log_file

# 写入日志
def write_log(log_file, message):
    with open(log_file, 'a') as f:
        f.write(message + "\n")

# 自定义随机森林分类器
class MyRandomForestClassifier:
    def __init__(self, tree_num=10, max_depth=None, log_file=None):
        self.tree_num = tree_num
        self.max_depth = max_depth
        self.trees = [DecisionTreeClassifier(max_depth=max_depth) for _ in range(tree_num)]
        self.log_file = log_file

    def fit(self, x_train, y_train):
        subset_size = len(x_train)
        for i in range(len(self.trees)):
            bootstrap_indices = np.random.choice(np.arange(x_train.shape[0]), size=subset_size, replace=True)
            x_bootstrap = x_train[bootstrap_indices]
            y_bootstrap = y_train[bootstrap_indices]
            msg = f"Bootstrap sample {i + 1} contains {len(np.unique(bootstrap_indices))} unique instances out of {subset_size}."
            print(msg)
            write_log(self.log_file, msg)
            
            msg = f"Training tree {i + 1}..."
            print(msg)
            write_log(self.log_file, msg)
            
            self.trees[i].fit(x_bootstrap, y_bootstrap)
            
            msg = f"Training tree {i + 1} completed."
            print(msg)
            write_log(self.log_file, msg)
            
            msg = '-' * 75  # 分隔线
            print(msg)
            write_log(self.log_file, msg)

    def predict(self, x_test):
        y_pred_score = {i: {0: 0, 1: 0, 2: 0} for i in range(len(x_test))}
        accuracys = []

        for i, tree in enumerate(self.trees):
            msg = f"Tree {i + 1} voting..."
            print(msg)
            write_log(self.log_file, msg)
            
            y_pred = tree.predict(x_test)
            for k in range(len(y_pred)):
                y_pred_score[k][y_pred[k]] += 1
                
            msg = f"Voting result: {y_pred}"
            print(msg)
            write_log(self.log_file, msg)
            
            msg = f"Updated total voting result: {y_pred_score}"
            print(msg)
            write_log(self.log_file, msg)
            
            accuracy = tree.score(x_test, y_test)
            accuracys.append(accuracy)
            
            msg = f'Tree {i + 1} accuracy: {accuracy:.2f}'
            print(msg)
            write_log(self.log_file, msg)
            
            msg = '-' * 75
            print(msg)
            write_log(self.log_file, msg)

        msg = f"Final voting results: {y_pred_score}"
        print(msg)
        write_log(self.log_file, msg)
        
        y_pred = np.array([np.argmax(list(y_pred_score[i].values())) for i in range(len(x_test))])
        
        msg = f"Final prediction result: {y_pred}"
        print(msg)
        write_log(self.log_file, msg)
        
        return y_pred

    def score(self, x_test, y_test):
        y_pred = self.predict(x_test)
        accuracy = np.mean(y_pred == y_test)
        return accuracy

# 创建日志文件
log_file = create_log_file()

# 加载数据集
data = load_wine()
X, y = data.data, data.target

# 划分训练集和测试集
x_train, x_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

# 决策树模型
dc = DecisionTreeClassifier()
dc.fit(x_train, y_train)
y_pred_dc = dc.predict(x_test)
dc_accuracy = dc.score(x_test, y_test)

# 自定义随机森林模型
myrfc = MyRandomForestClassifier(tree_num=20, max_depth=5, log_file=log_file)
myrfc.fit(x_train, y_train)
y_pred_myrfc = myrfc.predict(x_test)
my_accuracy = myrfc.score(x_test, y_test)

# sklearn 随机森林模型
rfc = RandomForestClassifier(n_estimators=20, random_state=42)
rfc.fit(x_train, y_train)
y_pred_rfc = rfc.predict(x_test)
rfc_accuracy = rfc.score(x_test, y_test)

# 输出结果
msg = '-' * 75
print(msg)
write_log(log_file, msg)

msg = f'模型结果对比'
print(msg)
write_log(log_file, msg)

msg = '-' * 75
print(msg)
write_log(log_file, msg)

msg = f'决策树模型准确率: {dc_accuracy:.2f}'
print(msg)
write_log(log_file, msg)

msg = f'自定义随机森林模型准确率: {my_accuracy:.2f}'
print(msg)
write_log(log_file, msg)

msg = f'sklearn 随机森林模型准确率: {rfc_accuracy:.2f}'
print(msg)
write_log(log_file, msg)

# 写入最终日志信息
msg = "\n" + "=" * 75
msg += f"\n实验完成！日志文件已保存至: {log_file}\n"
msg += "=" * 75
print(msg)
write_log(log_file, msg)